from nets.swin_transformer import swin_encoder
from tensorflow.keras.layers import (UpSampling2D,Conv2D,Concatenate,
                                    BatchNormalization,ReLU,Add,)
from tensorflow.keras import Model


def swin_seg_model(input_shape,num_classes):
    inputs,feat2,feat3,feat4,feat5 = swin_encoder(input_shape=input_shape)
    # print(inputs.shape,feat2.shape,feat3.shape,feat4.shape,feat5.shape)
    # (None, 512, 512, 1) (None, 128, 128, 96) (None, 64, 64, 192) (None, 32, 32, 384) (None, 16, 16, 768) 
    
    # -------------------------------
    # 分割头, feature paramid network
    # -------------------------------
    # 上采样
    feat5 = Conv2D(256,1)(feat5)
    feat5 = BatchNormalization()(feat5)
    feat5 = ReLU()(feat5)
    p5_up = UpSampling2D(interpolation='bilinear')(feat5)
    # print(p5_up.shape) # (None, 32, 32, 256)
    # input('zz')

    # -----------------------------
    # 与feat4的融合
    # -----------------------------
    feat4 = Conv2D(256,1)(feat4)
    feat4 = BatchNormalization()(feat4)
    feat4 = ReLU()(feat4)

    f4p5 = Add()([p5_up,feat4])
    # print(f4p5.shape) #(None, 32, 32, 256)
    # input('zz')

    # ----------------------------
    # 与feat3的融合
    # ----------------------------
    feat3 = Conv2D(256,1)(feat3)
    feat3= BatchNormalization()(feat3)
    feat3 = ReLU()(feat3)
    # 上采样
    p4_up = UpSampling2D(interpolation='bilinear')(f4p5)

    f3p4 = Add()([feat3,p4_up])
    # print(f3p4.shape) # (None, 64, 64, 256)
    # input('zz')

    # ----------------------------
    # 与feat2的融合
    # ----------------------------
    feat2 = Conv2D(256,1)(feat2)
    feat2= BatchNormalization()(feat2)
    feat2 = ReLU()(feat2)
    # 上采样
    p3_up = UpSampling2D(interpolation='bilinear')(f3p4)

    f2p3 = Add()([feat2,p3_up])
    # print(f2p3.shape) # (None, 128, 128, 256)

    f1 = UpSampling2D((8,8),interpolation='bilinear')(feat5)
    f2 = UpSampling2D((4,4),interpolation='bilinear')(f4p5)
    f3 = UpSampling2D((2,2),interpolation='bilinear')(f3p4)
    f4 = f2p3

    out = Concatenate()([f1,f2,f3,f4])
    # print(out.shape) # (None, 128, 128, 1024)

    out = Conv2D(512,3,padding='same')(out)
    out = BatchNormalization()(out)
    out = ReLU()(out)

    pred = Conv2D(num_classes,1,activation='sigmoid')(out)
    pred = UpSampling2D((4,4),interpolation='bilinear')(pred)
    # print(pred.shape)# (None, 512, 512, 2)

    model = Model(inputs,pred,name='swin_seg')
    return model



# if __name__ == '__main__':
#     input_shape = [512,512,3]
#     num_classes = 2
    
#     model = swin_seg_model(input_shape,num_classes)
#     model.summary()

    
   